# Sleep Project - Pedro Bessone, Gautam Rao, Heather Schofield, Frank Schilbach, and Mattie Toma
# Purpose: to replicate Figure 6 for the main paper (the summary bar graph)
# Last updated: Feb 21, 2021

#--------------SETUP---------------

# load packages
if (!require("pacman")) install.packages("pacman")
pacman::p_load(tidyverse, ggthemes, glue, here)

#------------PREPROCESS DATA----------------
# Loading dataset
data <- read_delim("Datasets/summary_coefficients.csv", delim = ";") %>%
  filter(!is.na(coefficients_ns)) %>%
  mutate(coefficients_nap = as.numeric(gsub("*", "", coefficients_nap, fixed = T)),
         coefficients_ns = as.numeric(gsub("*", "", coefficients_ns, fixed = T)),
         se_nap = as.numeric(sub(" ", "", se_nap)),
         outcome = str_to_title(outcome))

# reorder so that overall is first 
data <- bind_rows(data[15,], data[1:14,])

# reshape long 
long <- data %>% 
  pivot_longer(starts_with("coefficients_"), names_to = "treatment", values_to = "coefs") %>%
  pivot_longer(starts_with("se_"), values_to = "se") %>%
  mutate(treatment = sub("coefficients_", "", treatment),
         name = sub("se_", "", name),
         UB = coefs + 1.645 * se,
         LB = coefs - 1.645 * se) %>%
  filter(treatment == name) %>%
  select(-name) %>%
  mutate(treatment = ifelse(treatment == "nap", "Nap Treatment", "Night-Sleep Treatment"),
         treatment = ordered(treatment, levels = c("Night-Sleep Treatment", "Nap Treatment"))) %>% 
  mutate(outcome = ifelse(outcome == "Time Preferences", "Time", outcome),
         outcome = ifelse(outcome == "Physical Health", "Health", outcome),
         outcome = ifelse(outcome == "Psych Health", "Psychological", outcome),
         outcome = ifelse(outcome == "Wellbeing Index", "Well-being Index", outcome))

# make categories
long$categories <- ordered(c(rep("Overall", 2), rep("Work", 8), rep("Well-being", 6), 
                             rep("Cognition", 6), rep("Preferences", 8)), 
                           levels = c("Overall", "Work", "Well-being", "Cognition", "Preferences"))

# make positions to place the bar graphs at
long$positions <- rep(c(5,  15, 20, 25, 30, 40, 45, 50, 60, 65, 70, 80, 85, 90, 95), each = 2)

# make positions to label the bar graphs at
label_pos <- c(5, 22, 45, 65, 87)
labels <- rep(str_to_upper(levels(long$categories)), each = 1)
outcome <- long$outcome

# make color palette
pal <- c("Work" = "darkgreen", "Well-being" = "maroon", 
         "Cognition" = "#0d98ba", "Preferences" = "#3e6180", "Overall" = "grey")

#------------------PLOT DATA--------------------

mytheme = theme_minimal() + 
  theme(plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "lines"),
        panel.border = element_blank(),
        panel.background = element_blank(),
        panel.grid = element_blank(),
        axis.text.x=element_text(angle=45,hjust=1, size = 8, color = "black", face = "bold"),
        axis.text.y=element_text(size = 8),
        legend.position = "none",
        axis.ticks = element_blank(),
        axis.title.y=element_text(size=11),
        strip.text = element_text(size = 9))

# Graph night-sleep

df.g = long %>% 
  mutate(outcome = ifelse(outcome %in% c("Well-being Index", "Cognitive Index", "Preferences"), "Index", outcome)) %>% 
  filter(treatment == "Night-Sleep Treatment")

ggplot(df.g, aes(x = positions, y = coefs, fill = categories)) + 
  mytheme +
  theme(axis.text.x = element_blank(),
        axis.title.y=element_text(hjust = -1.5)) +
  geom_hline( yintercept = 0, color = "black" ) +
  geom_col(width = 4, color = 'black') + 
  geom_errorbar(aes(ymin = LB, ymax = UB), width = 0.2, color = 'black', show.legend = F) + 
  scale_fill_manual(values = pal) +  
  scale_x_continuous(breaks = df.g$positions, labels = df.g$outcome) +
  scale_y_continuous(lim = c(-.2, .4), breaks = seq(-.1,.2,.1), expand = c(0,0)) +
  labs(y = "Effect Size (SDs)", x = "") + 
  annotate(geom = "text", label = "(a) Pooled Night-Sleep Treatments", x = 46, y = .25) -> p1; p1

# add labels
for (i in 1:5) {
  p1 <- p1 + annotate(geom = "text", x = label_pos[i], y = 0.35, label = labels[i], size = 3, fontface =2)
}

# Graph nap

df.g = long %>% 
  mutate(outcome = ifelse(outcome %in% c("Well-being Index", "Cognitive Index", "Preferences"), "Index", outcome)) %>% 
  mutate(outcome = ifelse(outcome %in% c("Attention"), "Work Task", outcome)) %>%
  mutate(outcome = ifelse(outcome %in% c("Cognitive Function"), "Lab Tasks", outcome)) %>%
  mutate(outcome = ifelse(outcome %in% c("Typing Time"), "Labor Supply", outcome)) %>% 
  filter(treatment == "Nap Treatment")

ggplot(df.g, aes(x = positions, y = coefs, fill = categories)) + 
  mytheme +
  geom_hline( yintercept = 0, color = "black" ) +
  geom_col(width = 4, color = 'black') + 
  geom_errorbar(aes(ymin = LB, ymax = UB), width = 0.2, color = 'black', show.legend = F) + 
  scale_fill_manual(values = pal) +  
  scale_x_continuous(breaks = df.g$positions, labels = df.g$outcome) +
  scale_y_continuous(lim = c(-.15,.4), breaks = seq(-.1,.2,.1), expand = c(0,0)) +
  labs(y = " ", x = "") + 
  annotate(geom = "text", label = "(b) Nap Treatment", x = 46, y = .38)-> p2

p = cowplot::plot_grid(p1,p2, ncol = 1, rel_heights = c(1,1.2)); p

## Save graphs
ggsave("Output/Main_paper/Figures/Figure6_summary_both.eps", p, width = 7.25, height = 5)
ggsave("Output/Main_paper/Figures/Figure6_summary_both.pdf", p, width = 7.25, height = 5)


